from pathlib import Path
import pandas as pd
import cn2an
from typing import Optional

# 验证依赖
try:
    import openpyxl
except ImportError:
    print("缺少依赖库openpyxl, 请执行: pip install openpyxl")
    raise

# 锅炉配置
GUOLU_DICT = {
    "管道": ["圆管", "钢管", "管道"],
    "型钢": ["圆钢", "工字钢", "扁钢", "槽钢", "角钢"],
    "零件": ["法兰", "大小头", "弯管", "弯头", "弯头(无直段)", "弯头（无直段）", "管接头", "异径管"],
    "部件": ["分离器", "换向装置", "换向阀", "三通", "三通阀", "截止阀", "风箱", "取样装置", "分配器", "可调缩孔", "混合器", "测粉装置"],
    "锁气器": ["锁气器"],
    "门类": ["风门", "隔离门", "调节门", "隔绝门", "插板门", "煤阐门", "防爆门"],
    "孔类": ["煤粉吹扫孔", "捅煤孔", "人孔", "手孔", "除灰孔"],
    "补偿器": ["补偿器", "补偿器(全高)", "补偿器(CE波形)", "补偿器(UV波形)"],
    "传动装置及执行机构": ["传动装置", "执行机构"],
    "支吊架": ["管夹", "管担", "管座", "支座", "支架套筒", "管卡", "横担", "吊板", "支架翼板", "螺母", "螺丝", "拉杆", "螺栓", 
            "垫圈", "垫板", "连接板", "吊杆座", "加强板", "补强板", "吊板", "自调平座", "底板", "阻尼器", "减震器", "拉撑杆"],
    "钢板": ["钢板"]
}

# 气机配置
QIJI_DICT = {
    "弯头": ["弯头", "弯头(无直段)", "弯头（无直段）"],
    "管夹": ["管夹"],
    "螺栓螺母": ["螺栓", "螺母"],
    "型钢": ["圆钢", "工字钢", "扁钢", "槽钢", "角钢"],
    "钢管": ["钢管"],
    "支座": ["支座"],
    "封头": ["封头"],
    "接管座": ["接管座"],
    "阀": ["阀"],
    "三通": ["三通"],
    "钢板": ["钢板"],
    "异径管": ["异径管"],
    "流量测量装置": ["流量测量装置"],
    "窥油管": ["窥油管"],
    "节流组件": ["节流组件"],
    "过滤器": ["过滤器"],
    "堵头": ["堵头"]
}

FORMAT_TYPE = 0 # 0: 锅炉配置 1: 气机配置

def read_and_clean(input_path: str) -> pd.DataFrame:
    """读取并清洗单个Excel文件"""
    try:
        df = pd.read_excel(
            input_path,
            sheet_name=3,
            header=4,
            skipfooter=1
        )
    except ValueError as e:
        print(f"工作表索引错误: {input_path} - {str(e)}")
        return pd.DataFrame()
    except Exception as e:
        print(f"读取文件失败: {input_path} - {str(e)}")
        return pd.DataFrame()
    
    return clean_dataframe(df)

def clean_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """数据清洗逻辑"""
    if df.empty:
        return df

    try:
        # 重命名列
        df.columns = [
            '序号', '图号或代号', '名称及规格', '单位', '数量', 
            '材料', '单重', '总重', '备注'
        ]
    except Exception as e:
        print(f"列重命名失败: {str(e)}")
        return df

    # 拆分名称及规格
    try:
        split_df = df['名称及规格'].str.split(r'\s+', n=1, expand=True)
        df[['名称', '型号及规范']] = split_df.fillna('/')
    except Exception as e:
        print(f"名称规格拆分失败: {str(e)}")
        df[['名称', '型号及规范']] = df['名称及规格'].apply(lambda x: [x, '/'])

    # 处理显示格式
    for col in ['材料', '数量', '单位', '单重', '总重', '型号及规范']:
        df[col] = df[col].apply(lambda x: x if pd.notnull(x) else '/')

    # 移除临时列
    try:
        df.drop(columns=['序号', '图号或代号', '名称及规格'], inplace=True)
    except KeyError:
        print("临时列删除失败，可能已不存在")

    return df

def merge_materials(df: pd.DataFrame) -> pd.DataFrame:
    """合并相同物料"""
    if df.empty:
        return df

    # 类型转换
    for col in ['单重', '数量', '总重']:
        df[col] = pd.to_numeric(df[col], errors='coerce')

    group_cols = ['名称', '型号及规范', '材料', '单重']
    agg_rules = {
        '数量': 'sum',
        '总重': 'sum',
        '备注': lambda x: ';'.join(x.dropna().unique()),
        '单位': 'first',
        '单重': 'first'
    }

    try:
        merged = df.groupby(group_cols, as_index=False).agg(agg_rules)
    except Exception as e:
        print(f"合并物料失败: {str(e)}")
        return df

    # 总重校验
    merged['总重'] = merged.apply(
        lambda row: row['单重'] * row['数量'] 
        if (pd.notnull(row['单重']) and pd.notnull(row['数量'])) 
        else row['总重'],
        axis=1
    )

    return merged

def merge_number(df: pd.DataFrame) -> pd.DataFrame:
    """合并单重为空的物料"""
    if df.empty:
        return df

    df['数量'] = pd.to_numeric(df['数量'], errors='coerce')

    group_cols = ['名称', '型号及规范', '材料', '单位', '单重', '总重']
    agg_rules = {
        '数量': 'sum',
        '备注': lambda x: ';'.join(x.dropna().unique()),
    }

    try:
        return df.groupby(group_cols, as_index=False).agg(agg_rules)
    except Exception as e:
        print(f"合并数量失败: {str(e)}")
        return df

def sort_and_group_materials(df: pd.DataFrame, sort_dict: dict) -> pd.DataFrame:
    """分类排序并添加序号"""
    if df.empty:
        return df

    df = df.copy()
    df['分类'] = None
    df.insert(0, '序号', None)

    # 分类映射
    for category, keywords in sort_dict.items():
        mask = df['名称'].str.endswith(tuple(keywords), na=False)
        df.loc[mask, '分类'] = category

    # 添加子类别列
    def get_subcategory(row):
        """根据名称匹配子类别"""
        name = row['名称']
        for category, keywords in sort_dict.items():
            for keyword in keywords:
                if name.endswith(keyword):
                    return keyword
        return None

    df['子类别'] = df.apply(get_subcategory, axis=1)

    # 排序：先按分类，再按子类别
    try:
        df = df.sort_values(by=['分类', '子类别'], na_position='last').reset_index(drop=True)
        df['序号'] = df.groupby('分类', dropna=False).cumcount() + 1
    except Exception as e:
        print(f"分类排序失败: {str(e)}")
        return df

    # 构建结果
    result = []
    for i, (category, group) in enumerate(df.groupby('分类', dropna=False), 1):
        try:
            header = cn2an.an2cn(i) # 转换为中文序号
        except Exception:
            header = str(i)
        
        header_row = pd.DataFrame([{
            '序号': header,
            '名称': category if not pd.isna(category) else '其他',
            '型号及规范': '',
            '材料': '',
            '单位': '',
            '数量': '',
            '单重': '',
            '总重': '',
            '备注': ''
        }])
        
        result.append(header_row)
        result.append(group)

    try:
        final_df = pd.concat(result, ignore_index=True)
    except Exception as e:
        print(f"合并分类结果失败: {str(e)}")
        return df

    return final_df.drop(columns=['分类', '子类别'])

def format_columns(df: pd.DataFrame) -> pd.DataFrame:
    """格式化输出列顺序"""
    column_order = [
        '名称', '型号及规范', '材料', '单位', 
        '数量', '单重', '总重', '备注'
    ]

    try:
        df = df.reindex(columns=column_order)
        if (FORMAT_TYPE == 1):
            return sort_and_group_materials(df, QIJI_DICT)
        else:
            return sort_and_group_materials(df, GUOLU_DICT)
    except Exception as e:
        print(f"格式化列失败: {str(e)}")
        return df

def save_to_excel(df: pd.DataFrame, output_path: str) -> None:
    """保存结果到Excel"""
    try:
        with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
            df.to_excel(writer, index=False)
    except Exception as e:
        print(f"保存文件失败: {output_path} - {str(e)}")
        raise

def process_all_data(combined_df: pd.DataFrame) -> pd.DataFrame:
    """处理完整数据集"""
    if combined_df.empty:
        return combined_df

    # 过滤无效数据
    combined_df = combined_df[combined_df['名称'] != '/']

    # 分离单重为空的数据
    combo_mask = combined_df['单重'].isna() | combined_df['单重'].isin(['', '-', '/'])
    combo_df = combined_df[combo_mask].copy()
    normal_df = combined_df[~combo_mask].copy()

    # 合并处理
    merged_df = merge_materials(normal_df)
    number_df = merge_number(combo_df)

    try:
        final_df = pd.concat([merged_df, number_df], ignore_index=True)
    except Exception as e:
        print(f"合并最终结果失败: {str(e)}")
        return pd.DataFrame()

    return format_columns(final_df)

def process_directory(input_dir: str, output_dir: str, format_type: int = 0) -> Optional[Path]:
    """处理目录中的所有Excel文件"""
    input_path = Path(input_dir)
    output_path = Path(output_dir) / "combined.xlsx"

    # 设置格式类型
    global FORMAT_TYPE
    FORMAT_TYPE = format_type

    try:
        output_path.parent.mkdir(parents=True, exist_ok=True)
    except Exception as e:
        print(f"创建输出目录失败: {str(e)}")
        return None

    # 收集文件
    excel_files = list(input_path.glob("*.xlsx")) + list(input_path.glob("*.xls"))
    if not excel_files:
        print("未找到Excel文件")
        return None

    all_dfs = []
    for file in excel_files:
        if not file.name.startswith('~$'):
            print(f"处理文件: {file.name}")
            df = read_and_clean(file)
            if not df.empty:
                all_dfs.append(df)

    if not all_dfs:
        print("所有文件处理后数据为空")
        return None

    try:
        combined_df = pd.concat(all_dfs, ignore_index=True)
    except Exception as e:
        print(f"合并数据失败: {str(e)}")
        return None

    final_df = process_all_data(combined_df)
    if final_df.empty:
        print("最终处理结果为空")
        return None

    try:
        save_to_excel(final_df, output_path)
        print(f"处理完成，结果保存至: {output_path}")
        return output_path
    except Exception as e:
        print(f"保存结果失败: {str(e)}")
        return None
